#include <stdlib.h>
#include <stdbool.h>
#include "AVL.h"

Node *createNode(int x) {
  Node *newNode = (Node*)malloc(sizeof(Node));
  newNode->key = x;
  newNode->height = 0;
  newNode->left = newNode->right = NULL;
  return newNode;
}

int getHeight(Node *root) {
  if(root == NULL) {
    return -1;
  } else {
    return root->height;
  }
}

// Updates the height of a node based on the heights of its subtrees (if any).
// See comment on slide 16, chapter 11.
// Note: this function assumes that root is not NULL.
void updateHeight(Node *root) {
  int leftHeight = getHeight(root->left);
  int rightHeight = getHeight(root->right);
  // Note: if root has no children then this correctly computes zero.
  root->height = (leftHeight > rightHeight ? leftHeight : rightHeight ) + 1;
}

// Checks whether the root is balanced or not.
// See slide 2, chapter 11.
// Note: this function assumes that root is not NULL.
bool balanced(Node *root) {
  int leftHeight = getHeight(root->left);
  int rightHeight = getHeight(root->right);
  int diff = leftHeight - rightHeight;
  int d = diff < 0 ? -diff : diff;
  if(d <= 1) {
    return true;
  } else {
    return false;
  }
}

// Computes the case (1, 2, 3, or 4) when the root is not balanced.
// See slide 2, chapter 11.
// Note: this function assumes that root is not NULL.
// Note: this function is only called when the tree is unbalanced (see the
// rebalance function below).  To be unbalanced, the tree must have at
// least one non-empty subtree.  If the tree has at least one non-empty
// subtree then the first "if" test in the code below guarantees that we
// will never look at an empty subtree (if any), which in turn guarantees
// that all our pointer dereferences root->...->left/right will never fail.
int checkCase(Node *root) {
  int leftHeight = getHeight(root->left);
  int rightHeight = getHeight(root->right);
  if(leftHeight > rightHeight) {
    // Tree is left heavy.
    int leftLeftHeight = getHeight(root->left->left);
    int leftRightHeight = getHeight(root->left->right);
    if(leftLeftHeight >= leftRightHeight) {
      // Tree is left left heavy.
      // The = above is to use left left (not left right) when the left
      // subtree is perfectly balanced (see slide 25, chapter 11).
      return 1;
    } else {
      // Tree is left right heavy.
      return 2;
    }
  } else {
    // Tree is right heavy.
    int rightLeftHeight = getHeight(root->right->left);
    int rightRightHeight = getHeight(root->right->right);
    if(rightRightHeight >= rightLeftHeight) {
      // Tree is right right heavy.
      // The = above is to use right right (not right left) when the right
      // subtree is perfectly balanced (see slide 25, chapter 11).
      return 4;
    } else {
      // Tree is right left heavy.
      return 3;
    }
  }
}

// See slides 3 and 18, chapter 11.
void rightRotate(Node **proot) {
  Node *k2 = *proot, *k1 = k2->left, *Y = k1->right;
  *proot = k1;
  k1->right = k2;
  k2->left = Y;
  updateHeight(k2);
  updateHeight(k1);
}

// See slide 4, chapter 11.
void leftRotate(Node **proot) {
  Node *k1 = *proot, *k2 = k1->right, *Y = k2->left;
  *proot = k2;
  k2->left = k1;
  k1->right = Y;
  updateHeight(k1);
  updateHeight(k2);
}

// See slide 17, chapter 11.
// Note: this function assumes that root is not NULL.
void rebalance(Node **proot) {
  if(balanced(*proot)) {
    return;
  }
  switch(checkCase(*proot)) {
  case 1: // left left
    rightRotate(proot);
    break;
  case 4: // right right
    leftRotate(proot);
    break;
  case 2: // left right
    leftRotate(&(*proot)->left);
    rightRotate(proot);
    break;
  case 3: // right left
    rightRotate(&(*proot)->right);
    leftRotate(proot);
    break;
  }
}

// See slide 16, chapter 11.
Node *insertNode(Node **proot, int x) {
  if(*proot == NULL) {
    return *proot = createNode(x);
  }
  if((*proot)->key == x) {
    return NULL;
  }
  Node *newNode;
  if((*proot)->key > x) {
    newNode = insertNode(&(*proot)->left, x);
  } else {
    newNode = insertNode(&(*proot)->right, x);
  }
  // Note: *proot can never be NULL here:
  updateHeight(*proot);
  rebalance(proot);
  return newNode;
}

// See slide 23, chapter 8.
Node *findNode(Node *root, int x) {
  if(root == NULL) {
    return NULL;
  }
  if(root->key == x) {
    return root;
  }
  if(root->key > x) {
    return findNode(root->left, x);
  } else {
    return findNode(root->right, x);
  }
}

// Deletes and returns as result the minimum node in the non-empty tree.
// Since the tree is non-empty, such minimum node must exist.
// See slide 34, chapter 11.
Node *deleteMin(Node **proot) {
  if((*proot)->left == NULL) {
    // root is the minimum node, remove it from the tree and return it
    Node *minNode = *proot;
    *proot = (*proot)->right;
    return minNode;
  } else {
    // Keep moving down-left
    Node *minNode =  deleteMin(&(*proot)->left);
    // Note: *proot can never be NULL here because minNode must have had a
    // parent node before it was removed from the tree:
    updateHeight(*proot);
    rebalance(proot);
    return minNode;
  }
}

// See slide 33, chapter 11.
Node *deleteRoot(Node **proot) {
  Node *currNode = *proot;
  // This works for both BST deletion case 1 (no children) and BST deletion
  // case 2-b (one right child):
  if((*proot)->left == NULL) {
    *proot = (*proot)->right;
    return currNode;
  }
  // BST deletion case 2-a (one left child):
  if((*proot)->right == NULL) {
    *proot = (*proot)->left;
    return currNode;
  }
  // BST deletion case 3 (two children):
  // Find the minimum node in the right subtree:
  Node *minNode = deleteMin(&(*proot)->right);
  // Replace the root with the minNode:
  minNode->left = (*proot)->left;
  minNode->right = (*proot)->right;
  *proot = minNode;
  // Return the deleted node (the old root):
  return currNode;
}

// See slide 32, chapter 11.
Node *deleteNode(Node **proot, int x){
  if(*proot == NULL) {
    // Cannot find key, deletion fails
    return NULL;
  }
  Node *matchNode;
  if((*proot)->key > x) {
    matchNode = deleteNode(&(*proot)->left, x);
  } else {
    if((*proot)->key < x) {
      matchNode = deleteNode(&(*proot)->right, x);
    } else {
      matchNode = deleteRoot(proot);
    }
  }
  // Note: at this point *proot might be NULL because we have just deleted
  // a node and the AVL tree might be empty now:
  if(*proot != NULL){
    updateHeight(*proot);
    rebalance(proot);
  }
  return matchNode;
}

void destroyTree(Node *root) {
  if(root == NULL) {
    return;
  }
  destroyTree(root->left);
  destroyTree(root->right);
  free(root);
}
